import os
import csv

def read_csv(file_dir=None, file_name="result.csv"):
    csv_dict = {}
    if file_dir is not None:
        file_name = os.path.join(file_dir, file_name)
    with open(file_name, "r", encoding="utf-8-sig") as csv_file:
        csv_content = csv.reader(csv_file)
        # header = next(csv_content)
        for row in csv_content:
            # print(row)
            csv_dict[row[0]] = row[1]
    return csv_dict

dst_dir = "/public/ai_platform/models/cmcc_2022/dataset_aishell/"

# wav file dict
def read_wavs(dst_dir):
    file_dict = {}
    for root, dirs, files in os.walk(dst_dir):
        for file in files:
            if file.endswith(".wav"):
                file_path = os.path.join(root, file)
                file_dict[file] = file_path
                print(file, file_path)
    return file_dict

# read ground truth
gt_dict = {}
gt_file = os.path.join(dst_dir, "transcript", "aishell_transcript_v0.8.txt")
gt_file = "/public/ai_platform/models/cmcc_2022/dataset_aishell/transcript/aishell_transcript_v0.8.txt"
with open(gt_file, "r", encoding="utf-8") as gt:
    line = gt.readline().strip()
    head = line.split(" ")[0]
    gt_dict[head] = line

# read sorted wav file list
# sorted_file = "val_map_sorted.txt"
sorted_file = "val_seq_len.txt"
test_dict = {}
with open(sorted_file, "r", encoding="utf-8") as sf:
    for line in sf.readlines():
        head = line.split("\t233\n")[0]
        head = head.strip("\t233\n").split("\t")[0]
        num = line[-5:]
        test_dict[head] = int(num)

print(test_dict)

def get_full_path(val_dict, local_wavs):
    for k, v in val_dict.items():
        if k in local_wavs.keys():
            val_dict[k] = local_wavs[k]
            print(k, local_wavs[k])

# 
writeto_file = "val_map_gt.txt"
def packing_to_txt(writeto_file):
    packing = []
    amount = 0
    line = []
    lines = []
    length = []
    lens = []
    average = []
    write_list = []
    
    for kw, v in test_dict.items():
        k = kw.split(".wav")[0]        

        line.append(kw)
        length.append(v)
        amount += v
        
        if amount > 2048: # 500
            line.remove(kw)
            length.remove(v)
            amount -= v
            
            average.append(int(sum(length)/len(length)))
            write_list.append(" ".join(line) + f" {amount} {int(sum(length)/len(length))}\n")
            
            packing.append(" ".join(line))
            lines.append(line)
            lens.append(length)
            
            
            amount = 0
            line = []
            line.append(kw)
            length = []
            length.append(v)
            amount += v
    
    # for n, i in enumerate(average):
        
    
    
    # with open(writeto_file, "a+", encoding="utf-8") as val_gt:
    #     pass
    #     val_gt.write()


    print(average)
# packing_to_txt(writeto_file)

if __name__ == "__main__":
    
    local_wavs = "local_wavs.txt"
    if not os.path.exists(local_wavs):
        wavs = read_wavs(dst_dir)
        with open(local_wavs, "a+", encoding="utf-8") as ws:
            for k, v in wavs.items():
                ws.write(k + " " + v + "\n")

    wavs_dict = {}
    with open(local_wavs, "r") as wavs:
        for line in wavs.readlines():
            file_name = line.strip("\n").split(" ")[0]
            wavs_dict[file_name] = line.split(" ")[1]
    
    # get_full_path(test_dict, wavs_dict)
    
    # full_path = "full_path.txt"
    # with open(full_path, "a+") as fp:
    #     for k, v in test_dict.items():
    #         fp.write(k + " " + v)
    
    data_compare = 1
    if data_compare:
        test_csv = "/home/gyf/pkg/conformer_new/asr_conformer/tensorrt/result_test.csv"
        test_dict = read_csv(file_name=test_csv)
        
        
        pack_csv = "/home/gyf/pkg/conformer_new/asr_conformer/tensorrt/result_test_pack_onnx.csv"
        pack_list = "/home/gyf/pkg/conformer_new/asr_conformer/dataset/val_map_gt.txt"
        with open(pack_list, "r", encoding="utf-8") as pl, open(pack_csv, "a+", encoding="utf-8") as pc:
            lines = pl.readlines()
            
            writer = csv.writer(pc)
                    
            for line in lines:
                pack = []
                line = line.strip().split(" ")
                line.pop(-1)
                line.pop(-1)
                for wav in line:
                    pack.append(test_dict[wav])
                
                writer.writerow((line[0], "".join(pack)))
                